from khronos.opt.mh.treesearch.sts import STS, NodeWrapper, move_score

class FastSTS(STS):
    """Base class implementing the *fast* stochastic tree search algorithm for combinatorial 
    optimization problems. See the documentation of the base STS class for details on how to
    use STS and FastSTS algorithms.
    The fast version only requires one more method to be defined, which is optional but advisable.
        child_rank(parent, move, parent_rank)  # Compute the rank of a child node.
    
    IMPORTANT NOTE: The fast stochastic tree search algorithm is a slightly modified version of 
    the original STS algorithm. The intent of the modifications to STS is to enhance the speed of 
    the search, and reduce the computational costs of evaluating nodes before they are actually 
    explored. This is achieved by only adding to the queue the information necessary to create 
    each node (parent node + move) instead of creating a copy of the parent node and applying 
    the move to the copy before queueing. This way, the number of (potentially expensive) copy 
    operations on nodes is greatly reduced. Also, on the original STS, several of these copies 
    are dropped immediately after they are created, either because of pruning or queue discarding.
    """
    
    def bootstrap(self):
        """Create the root node's children and queue them before the main loop starts. If the 
        root is a leaf, its solution is checked, and the no node is queued. The root is pruned, 
        and also not queued if its lower bound is >= than the initial upper bound."""
        root = self.root()
        lower_bound = self.lower_bound(root)
        self.optimal = max(self.optimal, lower_bound)
        if self.is_leaf(root):
            self.check_leaf(root)
        elif lower_bound >= self.upper_bound:
            self.meta.Pruned += 1
        else:
            root_wrapper = NodeWrapper(root, lower_bound, None)
            self.queue_children(root_wrapper, None)
            
    def fetch_next(self):
        """Get the next node to explore from the queue, generate it (only copy parent and apply 
        move at this time), and return it to be explored. If the newly created node is a leaf, 
        its solution is checked, and a new node must be fetched. Also, if the node's lower bound 
        is >= than the current upper bound, the node is pruned and a new node must be fetched. 
        If the node passes through all these checks then it is finally returned."""
        position = self.select_next()
        index = int(position * (len(self.queue) - 1) + 0.5)
        current = self.queue.pop_at(index)
        node = self.copy(current.parent.node)
        self.apply(node, current.move)
        # Check the modified node's "leafness" and lower bound.
        if self.is_leaf(node):
            self.check_leaf(node)
            return None
        lower_bound = self.lower_bound(node)
        if lower_bound >= self.upper_bound:
            self.pruned += 1
            return None
        return NodeWrapper(node, self.rank(node), lower_bound)
        
    def expand(self):
        """Expand the current node by creating a child node for each move available at the 
        current node, and attempt to place the child nodes in the queue.
        NOTE: This method is used only if the 'diving' flag is False at initialization."""
        self.queue_children(self.current, None)
        self.explored += 1
        
    def dive(self):
        """Make a dive from the current node, descending through a greedy path and adding to the 
        queue any non-greedy child nodes produced while descending. Diving allows us to quickly 
        obtain solutions, but the quality of these early solutions depends heavily on the 
        performance of the branching heuristic (the moves() method). If during the dive, the 
        greedy node's lower bound is greater or equal to the current upper bound, the dive is 
        aborted and that node is pruned.
        NOTE: This method is used only if the 'diving' flag is True at initialization."""
        node = self.current.node
        rank = self.current.rank
        lower_bound = self.current.lower_bound
        while True:
            scored_moves = self.moves(node)
            greedy_path = self.dive_path(node, scored_moves)
            scored_moves.remove(greedy_path)
            # Attempt to queue non-greedy child nodes.
            if len(scored_moves) > 0:
                scored_moves.sort(key=move_score)
                self.queue_children(NodeWrapper(node, rank, lower_bound), scored_moves)
            # Apply the greedy move to a copy of the current node, and count the node 
            # as explored. Afterwards, check if the new node is a leaf, check its 
            # lower bound, and compute its rank for the next iteration of the dive.
            greedy_move, greedy_score = greedy_path
            node = self.copy(node)
            self.apply(node, greedy_move)
            self.explored += 1
            if self.is_leaf(node):
                self.check_leaf(node)
                break
            lower_bound = self.lower_bound(node)
            if lower_bound >= self.upper_bound:
                self.pruned += 1
                break
            rank = self.rank(node)
        self.dives += 1
        
    def queue_children(self, parent, scored_moves):
        """Attempt to add a set of child nodes (one per move) to the queue."""
        parent.rank = (parent.rank, self.rng.random())
        index = self.queue.bisect_left(parent)
        if scored_moves is None:
            scored_moves = self.moves(parent.node)
            scored_moves.sort(key=move_score)
        r = len(self.queue) - index
        self.queue.rotate(r)
        self.queue.extend([ChildWrapper(parent, move, score) for move, score in scored_moves])
        self.queue.rotate(-r)
        while len(self.queue) > self.queue_limit:
            self.queue.popleft()
            self.discarded += 1
            
    def child_rank(self, parent, move):
        """Compute and return the rank of a child node based on information of its parent node 
        and move. The rank of the child node is of extremme importance as it helps determine 
        the search order."""
        raise NotImplementedError()
        
class ParentWrapper(object):
    """Parent node wrapper for *fast* stochastic tree search. It is used to hold the node's 
    state and lower bound.
    Attributes:
        node - The node's state information.
        lower_bound - A lower bound of the best z-value obtained through this node.
    """
    __slots__ = ["node", "lower_bound"]
    
    def __init__(self, node, lower_bound):
        self.node = node
        self.lower_bound = lower_bound
        
class ChildWrapper(object):
    """Child node wrapper for *fast* stochastic tree search. Contains only a reference to the 
    parent node's wrapper, move data necessary to generate the actual child node, and the child's 
    rank. Also, the child node contains a lower_bound (used in queue pruning) descriptor that 
    refers to the respective data in the parent's wrapper."""
    __slots__ = ["parent", "move", "rank"]
    
    def __init__(self, parent, move, rank):
        self.parent = parent
        self.move = move
        self.rank = rank
        
    def __repr__(self):
        return "<Child(P=%s, R=%s, LB=%s) at %s>" % (hex(id(self.parent)), self.rank, 
                                                     self.parent.lower_bound, hex(id(self)))
        
    def __cmp__(self, node):
        return cmp(self.rank, node.rank)
        
    def lower_bound(self):
        return self.parent.lower_bound
        
    lower_bound = property(lower_bound)
    
